#include <iostream>
#include <chrono>
#include <random>
#include <thread>
#include <mpi.h>

// 函数原型声明
double run_single_thread_test(int proc_count, int compute_iter, int outer_iter);
double cpu_intensive_task(int iterations, int seed);
double mpi_worker_compute(int rank, int compute_iter, int outer_idx);
void simulated_sync_operation(int rank, int seed);

// 主函数
int main(int argc, char**argv) {
    MPI_Init(&argc, &argv);
    
    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    
    // 默认8个进程
    int proc_count = 8;
    int compute_iter = 100000;
    int outer_iter = 10;  // 外部循环次数
    bool enable_single_thread = true;
    
    // 解析命令行参数
    if (argc > 1) proc_count = std::stoi(argv[1]);
    if (argc > 2) compute_iter = std::stoi(argv[2]);
    if (argc > 3) outer_iter = std::stoi(argv[3]);
    if (argc > 4) enable_single_thread = (std::stoi(argv[4]) != 0);
    
    // 主进程输出配置信息
    if (rank == 0) {
        std::cout << "===== MPI多进程CPU仿真器测试 =====" << std::endl;
        std::cout << "进程数量: " << proc_count << std::endl;
        std::cout << "每个进程计算循环次数: " << compute_iter << std::endl;
        std::cout << "外部循环次数: " << outer_iter << std::endl;
        std::cout << "工作进程仅执行计算，同步由主进程协调" << std::endl;
        std::cout << "==============================" << std::endl;
    }
    
    // 单线程测试（仅主进程执行）
    double single_thread_duration = 0.0;
    if (rank == 0 && enable_single_thread) {
        std::cout << "开始单线程测试..." << std::endl;
        single_thread_duration = run_single_thread_test(proc_count, compute_iter, outer_iter);
        std::cout << "单线程总耗时: " << single_thread_duration << " 秒" << std::endl;
    }
    
    // 所有进程同步，确保单线程测试完成后再开始MPI测试
    MPI_Barrier(MPI_COMM_WORLD);
    
    double start_time, end_time;
    if (rank == 0) {
        std::cout << "开始MPI多进程测试..." << std::endl;
        start_time = MPI_Wtime();
    }
    
    // 外部循环 - 每次循环都执行计算和同步
    for (int outer_idx = 0; outer_idx < outer_iter; ++outer_idx) {
        // 主进程每5次循环输出一次进度
        if (rank == 0 && outer_idx % 5 == 0) {
            std::cout << "执行外部循环 " << outer_idx + 1 << "/" << outer_iter << std::endl;
        }
        
        // 1. 工作进程仅执行计算任务，无内部同步逻辑
        double local_result = mpi_worker_compute(rank, compute_iter, outer_idx);
        
        // 2. 所有进程同步点（由MPI库协调）
        MPI_Barrier(MPI_COMM_WORLD);
        
        // 3. 主进程执行同步操作（汇总结果和模拟同步延迟）
        if (rank == 0) {
            // 汇总所有进程的计算结果
            double global_sum = local_result;  // 主进程自己的结果
            double temp_result;
            
            for (int i = 1; i < size; ++i) {
                MPI_Recv(&temp_result, 1, MPI_DOUBLE, i, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
                global_sum += temp_result;
            }
            
            // 执行同步操作（仅主进程）
            simulated_sync_operation(0, outer_idx * 10000);
        }
        // 非主进程发送计算结果给主进程
        else {
            MPI_Send(&local_result, 1, MPI_DOUBLE, 0, 0, MPI_COMM_WORLD);
        }
        
        // 确保主进程完成同步操作后，所有进程再进入下一次循环
        MPI_Barrier(MPI_COMM_WORLD);
    }
    
    // 计算总耗时并输出结果
    if (rank == 0) {
        end_time = MPI_Wtime();
        double mpi_duration = end_time - start_time;
        std::cout << "多进程总耗时: " << mpi_duration << " 秒" << std::endl;
        
        if (enable_single_thread) {
            double speedup = single_thread_duration / mpi_duration;
            std::cout << "加速比: " << speedup << std::endl;
        }
        
        std::cout << "===== 结果汇总 =====" << std::endl;
        std::cout << "单线程耗时: " << single_thread_duration << std::endl;
        std::cout << "多进程耗时: " << mpi_duration << std::endl;
        if (enable_single_thread) {
            std::cout << "加速比: " << (single_thread_duration / mpi_duration) << std::endl;
        }
    }
    
    MPI_Finalize();
    return 0;
}

// MPI工作进程计算函数 - 仅执行计算任务，不包含同步逻辑
double mpi_worker_compute(int rank, int compute_iter, int outer_idx) {
    // 仅执行计算任务，无任何同步操作
    return cpu_intensive_task(compute_iter, outer_idx * 1000 + rank);
}

// 单线程测试实现 - 仅执行计算任务
double run_single_thread_test(int proc_count, int compute_iter, int outer_iter) {
    auto start_time = std::chrono::high_resolution_clock::now();
    
    // 单线程模拟：仅执行计算任务
    for (int outer_idx = 0; outer_idx < outer_iter; ++outer_idx) {
        for (int i = 0; i < proc_count; ++i) {
            // 仅执行计算任务
            double local_result = cpu_intensive_task(compute_iter, outer_idx * 1000 + i);
        }
    }
    
    auto end_time = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> total_time = end_time - start_time;
    return total_time.count();
}

// 模拟同步操作（仅主进程执行）
void simulated_sync_operation(int rank, int seed) {
    std::mt19937 rng(seed);
    
    // 准备阶段延迟（1-10ms）
    std::uniform_real_distribution<double> prepare_delay(0.001, 0.01);
    double prepare_time = prepare_delay(rng);
    std::this_thread::sleep_for(std::chrono::duration<double>(prepare_time));
    
    // 传输/确认阶段延迟（2-15ms）
    std::uniform_real_distribution<double> transfer_delay(0.002, 0.015);
    double transfer_time = transfer_delay(rng);
    
    // 主进程额外延迟
    if (rank == 0) {
        std::uniform_real_distribution<double> master_extra_delay(0.003, 0.02);
        transfer_time += master_extra_delay(rng);
    }
    
    std::this_thread::sleep_for(std::chrono::duration<double>(transfer_time));
    
    #ifdef DEBUG
    std::cout << "同步操作完成，耗时: " 
              << (prepare_time + transfer_time) * 1000 << "ms" << std::endl;
    #endif
}

// CPU密集型任务实现
double cpu_intensive_task(int iterations, int seed) {
    std::mt19937 rng(seed);
    std::uniform_real_distribution<double> dist(0.0001, 0.001);
    
    double result = 0.0;
    for (int i = 0; i < iterations; ++i) {
        result += std::sin(i * 0.001) * std::cos(i * 0.002);
        if (i % 10000 == 0) {
            std::this_thread::sleep_for(std::chrono::duration<double>(dist(rng)));
        }
    }
    return result;
}
    